<template>
    <div class="app-container">
        <el-card class="box-card">
            <div slot="header" class="clearfix">
                <span class="card-title">模型训练</span>
                <!-- 根据数据集名称和验证状态查询 -->
                <div style="margin-left: 15px;width: 350px;float: right;">
                    <el-input placeholder="搜索模型名称" v-model="listQuery.model_name" clearable>ss="input-with-select">
                        <el-select v-model="listQuery.model_type" slot="prepend" placeholder="请选择">
                            <el-option label="YOLOv5" value="yolo"></el-option>
                            <el-option label="SlowFast" value="slowfast"></el-option>
                        </el-select>
                        <el-button slot="append" icon="el-icon-search"></el-button>
                    </el-input>
                </div>
                <!-- 新建模型按钮 -->
                <el-button type="primary" plain icon="el-icon-plus" style="float: right" @click="showCreateDialog">
                    新建模型
                </el-button>
            </div>

            <!-- 模型表格 -->
            <el-table v-loading="listLoading" :data="modelsData" element-loading-text="Loading" border fit
                highlight-current-row>
                <el-table-column align="center" label="ID" width="95">
                    <template slot-scope="scope">
                        {{ scope.row.id }}
                    </template>
                </el-table-column>
                <el-table-column label="模型名称" width="200" align="center">
                    <template slot-scope="scope">
                        {{ scope.row.name }}
                    </template>
                </el-table-column>
                <el-table-column label="描述" width="200" align="center">
                    <template slot-scope="scope">
                        <span>{{ scope.row.description }}</span>
                    </template>
                </el-table-column>
                <el-table-column label="数据集" width="200" align="center">
                    <template slot-scope="scope">
                        {{ scope.row.dataset }}
                    </template>
                </el-table-column>
                <el-table-column v-if='currentTask' class-name="status-col" label="训练状态" width="110" align="center">
                    <template slot-scope="scope">
                        <el-tag :type="scope.row.status | statusFilter">{{ scope.row.status }}</el-tag>
                    </template>
                </el-table-column>
                <el-table-column align="center" prop="created_time" label="创建时间" width="300">
                    <template slot-scope="scope">
                        <i class="el-icon-time" />
                        <span>{{ scope.row.created_time }}</span>
                    </template>
                </el-table-column>
                <el-table-column align="center" label="操作">
                    <template slot-scope="scope">
                        <el-button @click="showDetail(scope.row)">详情</el-button>
                        <el-button @click="handleTrainModel(scope.row)">训练</el-button>
                        <el-button @click="handleEditModel(scope.row)">编辑</el-button>
                        <el-button @click="handleDeleteModel(scope.row)">删除</el-button>
                    </template>
                </el-table-column>
            </el-table>

            <!-- 分页 -->
            <el-pagination v-show="total > 0" :current-page="currentPage" :page-size="pageSize"
                layout="prev, pager, next" :total="total" @current-change="handlePageChange" style="margin-top: 20px">
            </el-pagination>
        </el-card>

        <!-- 新建模型表单 -->
        <el-dialog title="新建模型" :visible.sync="dialogFormVisible" width="500px">
            <el-form ref="modelForm" :model="modelForm" :rules="rules" size="medium" label-width="100px">
                <el-form-item label="模型名称" prop="name">
                    <el-input v-model="modelForm.name" placeholder="请输入模型名称" :maxlength="11" show-word-limit
                        clearable></el-input>
                </el-form-item>
                <el-form-item label="模型描述" prop="description">
                    <el-input v-model="modelForm.description" type="textarea" placeholder="请输入模型描述"
                        :autosize="{ minRows: 4, maxRows: 4 }"></el-input>
                </el-form-item>
                <el-form-item label="数据集" prop="dataset_id">
                    <el-select v-model="modelForm.dataset_id" placeholder="请选择数据集" clearable>
                        <el-option v-for="(item, index) in datasetOptions" :key="index" :label="item.label"
                            :value="item.value" :disabled="item.disabled"></el-option>
                    </el-select>
                </el-form-item>
                <el-form-item label="模型类型" prop="modelType">
                    <el-select v-model="modelForm.modelType" placeholder="请选择模型类型" clearable>
                        <!-- <el-option v-for="(item, index) in modelTypeOptions" :key="index" :label="item.label"
                            :value="item.value" :disabled="item.disabled"></el-option> -->
                        <el-option label="YOLOv5" value="yolo"></el-option>
                        <el-option label="SlowFast" value="slowfast"></el-option>
                    </el-select>
                </el-form-item>
            </el-form>
            <div slot="footer" class="dialog-footer">
                <el-button @click="dialogFormVisible = false">取 消</el-button>
                <el-button type="primary" @click="handleConfirm">确 定</el-button>
            </div>
        </el-dialog>



        <!-- 训练模型参数表单 -->
        <el-dialog title="训练模型" :visible.sync="trainModelFormVisible" width="500px">
            <el-row :gutter="15">
                <el-form ref="argumentForm" :model="argumentForm" :rules="trainRules" size="medium" label-width="150px">
                    <el-col :span="23">
                        <el-form-item label="模型名称" prop="name">
                            <el-input v-model="argumentForm.name" placeholder="模型名称" :disabled='true'
                                clearable></el-input>
                        </el-form-item>
                    </el-col>
                    <el-col :span="23">
                        <el-form-item label="当前数据集" prop="dataset">
                            <el-input v-model="argumentForm.dataset" placeholder="当前数据集" :disabled='true'></el-input>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="输入图像大小" prop="imgsz">
                            <el-input-number v-model="argumentForm.imgsz" placeholder="输入图像大小" step-strictly>
                            </el-input-number>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="输入通道数" prop="channel_num">
                            <el-input-number v-model="argumentForm.channel_num" placeholder="输入通道数"></el-input-number>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="类别数量" prop="nc">
                            <el-input-number v-model="argumentForm.nc" placeholder="类别数量"></el-input-number>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="训练周期" prop="epochs">
                            <el-input-number v-model="argumentForm.epochs" placeholder="训练周期"></el-input-number>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="批量大小" prop="batch_size">
                            <el-input-number v-model="argumentForm.batch_size" placeholder="批量大小"></el-input-number>
                        </el-form-item>
                    </el-col>
                    <el-col :span="24">
                        <el-form-item label="学习率" prop="lr0">
                            <el-input v-model="argumentForm.lr0" placeholder="请输入学习率"></el-input>
                        </el-form-item>
                    </el-col>
                </el-form>
            </el-row>
            <div slot="footer">
                <el-button @click="trainModelFormVisible = false">取消</el-button>
                <el-button type="primary" @click="startTrainModel">开始训练</el-button>
            </div>
        </el-dialog>
    </div>
</template>

<script>
import axios from 'axios';
axios.defaults.baseURL = 'http://localhost:5000'
import api from '@/api/model'

export default {
    name: 'ModelTrain',
    data() {
        return {
            // 搜索栏目
            listQuery: {
                model_name: '',
                model_type: '',
            },
            // 表格相关数据
            modelsData: [],
            datasetOptions: [],// 数据集列表
            listLoading: false,
            //分页
            currentPage: 1,
            pageSize: 10,
            total: 0,
            // 新建模型表单
            dialogFormVisible: false,
            modelForm: {
                name: '',
                description: '',
                dataset_id: '',
                model_type: ''
            },
            rules: {
                name: [
                    { required: true, message: '模型名称不能为空', trigger: 'blur' },
                    { max: 50, message: '长度不能超过50个字符', trigger: 'blur' }
                ],
                dataset_id: [
                    { required: true, message: '请选择数据集文件', trigger: 'change' }
                ]
            },
            // 训练模型参数表单
            trainModelFormVisible: false,
            argumentForm: {
                name: ' ',
                dataset: ' ',
                imgsz: 256,
                channel_num: 3,
                nc: 1,
                epochs: 50,
                batch_size: 8,
                lr0: 0.001,
            },
            // 训练参数验证规则
            trainRules: {
                epochs: [{ type: 'number', min: 1, message: '至少训练1个epoch', trigger: 'blur' }],
                lr0: [{ pattern: /^\d+\.?\d*$/, message: '请输入有效数字', trigger: 'blur' }]
            },
            // 训练任务状态
            currentTask: null
        };
    },
    mounted() {
        this.fetchModels();
        this.fetchDatasets();
    },
    methods: {
        // 分页方法
        handlePageChange(page) {
            this.currentPage = page;
            this.fetchModels();
        },
        // 获取模型列表方法
        fetchModels() {
            this.listLoading = true;
            axios.get('/models/info', {
                params: {
                    //分页
                    page: this.currentPage,
                    page_size: this.pageSize,
                    //查询参数
                    model_type: this.listQuery.model_type,
                    model_name: this.listQuery.model_name
                }
            }).then(res => {
                this.modelsData = res.data.data.items;
                this.total = res.data.data.total;
            }).catch(error => {
                this.$message.error('获取数据集失败: ' + error.message)
                console.error('API Error:', error.response || error)
            }).finally(() => {
                this.listLoading = false;
            });
        },
        // 获取数据集列表方法
        fetchDatasets() {
            axios.get('/datasets/info').then(res => {
                this.datasetOptions = res.data.data.items.map(item => ({
                    value: item.id,
                    label: item.name
                }));
            });
        },
        //点击新建模型按钮
        showCreateDialog() {
            this.dialogFormVisible = true;//打开弹窗
        },
        // 提交新建模型表单
        handleConfirm() {
            this.$refs.modelForm.validate(valid => {
                console.log(this.modelForm)
                if (valid) {
                    axios.post('/models/create', this.modelForm).then(res => {
                        this.dialogFormVisible = false;// 关闭弹窗
                        this.fetchModels();
                    });
                }
            });
        },
        // 训练模型
        handleTrainModel(row) {
            this.trainModelFormVisible = true; // 显示训练模型表单
            this.argumentForm.name = row.name; // 设置模型名称
            this.argumentForm.dataset = row.dataset; // 设置数据集名称
            this.argumentForm.dataset_id = row.dataset_id; // 设置数据集ID
            this.selectedModel = row.model_type; // 设置模型类型
        },
        // 开始训练模型
        startTrainModel() {
            axios.post('/model/train', {
                dataset_id: row.dataset_id, // 数据集ID
            }).then(res => {
                this.currentTask = { id: res.data.task_id, status: 'pending' };
                this.pollTaskStatus();
            });
        },
        pollTaskStatus() {
            const check = () => {
                axios.get(`/api/tasks/${this.currentTask.id}`).then(res => {
                    this.currentTask.status = res.data.status;
                    if (res.data.status === 'training') {
                        setTimeout(check, 2000);
                    }
                });
            };
            check();
        }
    }
};
</script>

<style>
.el-select .el-input {
    width: 100%;
}

.input-with-select .el-input-group__prepend {
    background-color: #fff;
}

.el-form .el-select {
    width: 100%;
}
</style>